# Copyright (c) OpenMMLab. All rights reserved.
import argparse
import json

from xtuner.dataset.refcoco_json import RefCOCOJsonDataset


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--ann-path",
        default="data/refcoco_annotations",
        help="Refcoco annotation path",
    )
    parser.add_argument(
        "--image-path",
        default="data/llava_data/llava_images/coco/train2017",
        help="COCO image path",
    )
    parser.add_argument(
        "--save-path", default="./", help="The folder to save converted data"
    )
    args = parser.parse_args()
    return args


if __name__ == "__main__":
    args = parse_args()

    data_info = [
        ("refcoco", "unc"),
        ("refcoco+", "unc"),
        ("refcocog", "umd"),
    ]
    all_data = []
    for dataset, split in data_info:
        data = RefCOCOJsonDataset.get_data_json(
            ann_path=args.ann_path,
            image_path=args.image_path,
            dataset=dataset,
            splitBy=split,
        )[0]
        all_data.extend(data)
    save_path = args.save_path + "/train.json"
    with open(save_path, "w") as f:
        print(f"save to {save_path} with {len(all_data)} items.")
        print(all_data[0])
        json.dump(all_data, f, indent=4)
